import re
import subprocess
from flashtext import KeywordProcessor
import sys
from pathlib import Path
import os
import shutil 

# 初始化关键字处理器
keyword_processor = KeywordProcessor()
keyword_processor.add_keyword('Unsafe', '')


def copy_project(src_path, dest_path):
    """
    复制项目到新目录
    Args:
        src_path: 源项目路径
        dest_path: 目标路径
    Returns:
        bool: 复制是否成功
    """
    try:
        if os.path.exists(dest_path):
            shutil.rmtree(dest_path)
        shutil.copytree(src_path, dest_path)
        return True
    except Exception as e:
        print(f"复制项目时发生错误: {str(e)}")
        return False

def remove_unsafe_sig(project_path):
    """移除所有的unsafe标记"""
    
    # 遍历项目目录下的所有.rs文件
    for rust_file in Path(project_path).rglob('*.rs'):
        # 读取文件内容
        with open(rust_file, 'r', encoding='utf-8') as f:
            content = f.read()
        
        # 使用正则表达式移除unsafe关键字
        # 移除函数声明中的unsafe
        modified_content = re.sub(r'unsafe\s+fn', 'fn', content)
        # 移除unsafe块
        modified_content = re.sub(r'unsafe\s*\{([^}]*)\}', r'\1', modified_content)
        
        # 写回文件
        with open(rust_file, 'w', encoding='utf-8') as f:
            f.write(modified_content)
            f.flush()

# 定义错误模式
error_pattern = r'error\[\w+\]:\s*'

def lint_rust_file(file_path):
    """
    处理 Rust 文件中的不安全代码
    Args:
        file_path: Rust 文件的路径
    """
    try:
        with open(file_path, 'r') as f:
            content = f.read()
            content = keyword_processor.replace_keywords(content)
            
            # 使用 clippy 命令进行代码质量检查
            output = subprocess.check_output(
                ['cargo', 'clippy', '--'], 
                input=content.encode(), 
                stderr=subprocess.PIPE
            )
            output = output.decode()
            print(output)
            # 提取出不安全的代码段
            unsafe_match = re.search(f'{error_pattern}(.+?)\[', output, re.DOTALL)

            if unsafe_match:
                unsafe_code = unsafe_match.group(1)
                # 判断不安全的代码是否是裸指针带来的不安全
                pointer_pattern = r'\*mut\s*|&mut\s*\*|Box::into_raw'
                if re.search(pointer_pattern, unsafe_code):
                    # 将裸指针用 Unsafe 封装成一个全新的变量
                    unsafe_code = re.sub(pointer_pattern, 'UnsafePointer::new', unsafe_code)
                    unsafe_code = f'unsafe {{ let mut _unsafe_pointer = {unsafe_code}; }}'
                else:
                    # 不是裸指针则直接用 Unsafe 进行包裹
                    unsafe_code = f'unsafe {{ {unsafe_code} }}'
                
                content = re.sub(error_pattern, unsafe_code, content)

            with open(file_path, 'w') as f:
                f.write(content)
                
    except Exception as e:
        print(f"处理文件 {file_path} 时发生错误: {str(e)}")
        print(e.stderr.decode())
        return False
    
    return True

def main(project_path=None):
    """
    主函数，遍历项目目录下所有的.rs文件并移除unsafe标记
    Args:
        project_path: 项目根目录的路径，默认为'src'
    """
    if project_path is None:
        project_path = 'src'
    
    if not os.path.exists(project_path):
        print(f"错误: {project_path} 目录不存在")
        sys.exit(1)

    # 创建新的项目目录路径
    new_project_path = os.path.join(os.path.dirname(project_path), "P5_unsafe_fix")

    # 复制项目
    print(f"正在将项目复制到 {new_project_path}...")
    if not copy_project(project_path, new_project_path):
        print("项目复制失败")
        sys.exit(1)
    
    try:
        print(f"开始移除 {new_project_path} 目录下的unsafe标记...")
        remove_unsafe_sig(new_project_path)
        print("unsafe标记移除完成")
        
        # 遍历新项目目录下的所有.rs文件进行后续处理
        success_count = 0
        fail_count = 0
        
        for rust_file in Path(new_project_path).rglob('*.rs'):
            print(f"正在处理文件: {rust_file}")
            if lint_rust_file(str(rust_file)):
                success_count += 1
            else:
                fail_count += 1
        
        print(f"\n处理完成:")
        print(f"成功处理文件数: {success_count}")
        print(f"处理失败文件数: {fail_count}")
        
        if fail_count > 0:
            raise Exception(f"处理失败的文件数: {fail_count}")
            
    except Exception as e:
        print(f"处理过程中发生错误: {str(e)}")
        print("正在删除处理失败的项目副本...")
        shutil.rmtree(new_project_path)
        sys.exit(1)
    
    print(f"处理成功！处理后的项目保存在: {new_project_path}")

if __name__ == "__main__":
    if len(sys.argv) > 1:
        main(sys.argv[1])
    else:
        main()